"""
MCP与大模型集成演示脚本
演示如何让大模型调用MCP服务
"""

import os
import sys
from typing import List
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, SystemMessage
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode
from typing_extensions import TypedDict
from typing import Annotated


# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


class State(TypedDict):
    messages: Annotated[list, add_messages]


# 初始化大模型
API_BASE_URL: str = os.getenv("API_BASE_URL", "https://api.deepseek.com")
API_KEY: str = os.getenv("API_KEY", "sk-878af5ef9666470d8f58425a15d498d8")
MODEL_NAME: str = os.getenv("MODEL_NAME", "deepseek-chat")

llm = init_chat_model(
    MODEL_NAME,
    base_url=API_BASE_URL,
    api_key=API_KEY
)


@tool
def mcp_tool(server: str, tool: str, parameters: dict):
    """调用MCP服务器上的工具"""
    # 实际实现中会调用MCP服务
    # 这里使用模拟结果
    return {
        "server": server,
        "tool": tool,
        "parameters": parameters,
        "result": f"这是来自 {server} 服务器上 {tool} 工具的模拟结果，参数为 {parameters}"
    }


@tool
def list_mcp_servers():
    """列出所有可用的MCP服务器"""
    # 实际实现中会从配置中获取服务器列表
    # 这里使用模拟数据
    return {
        "enabled": True,
        "servers": ["search-server", "calculator-server"],
        "server_configs": {
            "search-server": {"enabled": True, "url": "http://localhost:3000/search"},
            "calculator-server": {"enabled": True, "url": "http://localhost:3001/calc"}
        }
    }


@tool
def check_server_availability(server_name: str):
    """检查特定MCP服务器是否可用"""
    # 实际实现中会检查服务器状态
    # 这里使用模拟数据
    available_servers = ["search-server", "calculator-server"]
    return {
        "server": server_name,
        "available": server_name in available_servers
    }


@tool
def list_server_tools(server_name: str):
    """列出特定MCP服务器上可用的工具"""
    # 实际实现中会从服务器获取工具列表
    # 这里使用模拟数据
    server_tools = {
        "search-server": ["web_search", "news_search", "image_search"],
        "calculator-server": ["add", "subtract", "multiply", "divide"]
    }
    return {
        "server": server_name,
        "tools": server_tools.get(server_name, [])
    }


# 创建工具列表
tools = [mcp_tool, list_mcp_servers, check_server_availability, list_server_tools]
tool_node = ToolNode(tools)


def chatbot(state: State):
    """聊天机器人节点"""
    # 使用全局定义的LLM实例
    llm_with_tools = llm.bind_tools([mcp_tool, list_mcp_servers, check_server_availability, list_server_tools])
    message = llm_with_tools.invoke(state["messages"])
    return {"messages": [message]}


def route_tools(state: State):
    """
    工具路由函数
    根据消息是否包含工具调用来决定下一步
    """
    messages = state.get("messages", [])
    if messages:
        last_message = messages[-1]
        if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
            return "tools"
    return "__end__"


# 创建工作流图
graph_builder = StateGraph(State)

# 添加节点
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_node("tools", tool_node)

# 添加边
graph_builder.add_edge(START, "chatbot")
graph_builder.add_conditional_edges(
    "chatbot",
    route_tools,
    {"tools": "tools", "__end__": "__end__"}
)
graph_builder.add_edge("tools", "chatbot")
graph_builder.add_edge("chatbot", END)

# 编译工作流
graph = graph_builder.compile()


def run_demo():
    """运行演示"""
    print("=== MCP与大模型集成演示 ===")
    print("此演示展示大模型如何调用MCP服务")
    print()
    
    # 显示当前MCP配置（模拟）
    print("当前MCP配置:")
    print("  启用状态: True")
    print("  服务器列表: ['search-server', 'calculator-server']")
    print("  服务器配置: {'search-server': {...}, 'calculator-server': {...}}")
    print()
    
    # 检查search-server是否可用（模拟）
    print("Search服务器可用性: 是")
    print("Search服务器工具: ['web_search', 'news_search', 'image_search']")
    print()
    
    # 演示对话
    conversation_history = [
        SystemMessage(content="""你是一个能够调用MCP工具的AI助手。
你的任务是根据用户的问题，判断是否需要调用MCP工具来获取信息。
如果需要，你应该选择合适的工具并提供必要的参数。
特别地，你可以使用以下MCP相关工具：
1. list_mcp_servers - 列出所有MCP服务器
2. check_server_availability - 检查特定服务器是否可用
3. list_server_tools - 列出特定服务器上的工具
4. mcp_tool - 调用MCP服务器上的具体工具

当用户询问与搜索相关的问题时，你应该：
1. 首先检查search-server是否可用
2. 如果可用，列出其支持的工具
3. 调用web_search工具执行搜索"""),
        HumanMessage(content="你好，你能帮我搜索一下最新的AI技术发展吗？")
    ]
    
    print("开始对话演示:")
    print("用户: 你好，你能帮我搜索一下最新的AI技术发展吗？")
    
    # 运行工作流
    for step in graph.stream(
        {"messages": conversation_history},
        stream_mode="values"
    ):
        # 获取最新的消息
        messages = step.get("messages", [])
        if messages:
            last_message = messages[-1]
            if hasattr(last_message, 'content') and last_message.content:
                print(f"助手: {last_message.content}")
            elif hasattr(last_message, 'tool_calls') and last_message.tool_calls:
                print("助手正在调用工具:")
                for tool_call in last_message.tool_calls:
                    print(f"  - {tool_call['name']}: {tool_call['args']}")
    
    print()
    print("=== 演示结束 ===")


if __name__ == "__main__":
    run_demo()